# -*- coding: utf-8 -*-

"""剑指 Offer II 047. 二叉树剪枝
给定一个二叉树 根节点 root ，树的每个节点的值要么是 0，要么是 1。请剪除该二叉树中所有节点的值为 0 的子树。
节点 node 的子树为 node 本身，以及所有 node 的后代。

示例 1:
输入: [1,null,0,0,1]
输出: [1,null,0,null,1] 
解释:
只有红色节点满足条件“所有不包含 1 的子树”。
右图为返回的答案。

示例 2:
输入: [1,0,1,0,0,0,1]
输出: [1,null,1,null,1]
解释:

示例 3:
输入: [1,1,0,1,1,0,1,0]
输出: [1,1,0,1,1,null,1]
解释:

提示:
二叉树的节点个数的范围是 [1,200]
二叉树节点的值只会是 0 或 1"""


# Definition for a binary tree node.
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


class Solution:
    """深度优先搜索，对每个节点的操作是一致的，可利用递归实现，先写出递推过程，递归过程就能实现了"""
    def pruneTree(self, root: TreeNode) -> TreeNode:
        def cancut(node):
            if not (node.left or node.right):
                return node.val

            leftcut, rightcut = 0, 0
            if node.left:
                leftcut = cancut(node.left)
            if node.right:
                rightcut = cancut(node.right)

            if not leftcut:
                node.left = None
            if not rightcut:
                node.right = None
            
            if node.val or leftcut or rightcut:
                return 1
            else:
                return 0
        
        if not cancut(root):
            return None
        else:
            return root
